'''
* This is the projet for Brtc LlmOps Platform
* @Author Leon-liao <liaosiliang@alltman.com>
* @Description //TODO 
* @File: vector_database_service.py
* @Time: 2025/9/26
* @All Rights Reserve By Brtc
'''
from dataclasses import dataclass
from typing import Any
import weaviate
from flask import Flask
from flask_weaviate import FlaskWeaviate
from injector import inject
from langchain_core.documents import Document as LCDocument
from langchain_core.vectorstores import VectorStoreRetriever
from langchain_weaviate import WeaviateVectorStore
from weaviate.collections import Collection
from .embeddings_service import EmbeddingsService

# 向量数据库的数据集名称
COLLECTION_NAME = "Brtc_Llomps_Ai_Platform_DatasetWithOpenAI"

@inject
@dataclass
class VectorDatabaseService:
    """向量数据库服务"""
    weaviate:FlaskWeaviate
    embeddings_service:EmbeddingsService

    async def _get_client(self, flask_app:Flask):
        with flask_app.app_context():
            return weaviate.client


    @property
    def vector_store(self)->WeaviateVectorStore:
        return WeaviateVectorStore(
            client=self.weaviate.client,
            index_name=COLLECTION_NAME,
            text_key="text",
            embedding=self.embeddings_service.embeddings,
        )


    async def add_documents(self, documents:list[LCDocument], **keywords:Any):
        """往向量数据库中，新增文档， 将vector_store使用async 进行二次封装，避免在gevent中出现时间循环错误"""
        self.vector_store.add_documents(documents, **keywords)


    def get_retriever(self)->VectorStoreRetriever:
        """获取检索器"""
        return self.vector_store.as_retriever()


    @property
    def collection(self)->Collection:
        """获取向量数据库的数据集"""
        return self.weaviate.client.collections.get(COLLECTION_NAME)
